#include "global_planner/ompl_planner.h"

namespace Planner
{
    double ValidityChecker::clearance(const ob::State *state) const
    {
        const ob::RealVectorStateSpace::StateType *state3D =
            state->as<ob::RealVectorStateSpace::StateType>();

        double x = state3D->values[0];
        double y = state3D->values[1];
        double z = state3D->values[2];

        return calcDis(x, y, z) - secure_dis_;
    }

    double ValidityChecker::calcDis(const double &x, const double &y, const double &z) const
    {
        pcl::PointXYZ pt(x, y, z);
        std::vector<int> pointIdxNKNSearch(1);
        std::vector<float> pointNKNSquaredDistance(1);
        if (kdtree_->nearestKSearch(pt, 1, pointIdxNKNSearch, pointNKNSquaredDistance) > 0)
        {
            float squaredDis = pointNKNSquaredDistance[0];
            return sqrt(squaredDis);
        }
        std::cout << "kdtree search error\n";
        return -1;
    }

    bool OMPLPlanner::setStart(const Vector3d &pt)
    {
        double dis = checker_->calcDis(pt[0], pt[1], pt[2]);
        if (dis < secure_dis_)
        {
            if (print_info_)
                cout << "invaild start position, the obstacle gap is " << dis << endl;
            return false;
        }
        if (print_info_)
            cout << "set start position to: " << pt[0] << " " << pt[1] << " " << pt[2] << endl;
        start_pos_ = pt;
        return true;
    }

    bool OMPLPlanner::setGoal(const Vector3d &pt)
    {
        double dis = checker_->calcDis(pt[0], pt[1], pt[2]);
        if (dis < secure_dis_)
        {
            if (print_info_)
                cout << "invaild goal position, the obstacle gap is " << dis << endl;
            return false;
        }

        if (print_info_)
            cout << "set goal position to: " << pt[0] << " " << pt[1] << " " << pt[2] << endl;
        goal_pos_ = pt;
        return true;
    }

    void OMPLPlanner::setCloudMap(pcl::PointCloud<pcl::PointXYZ>::Ptr &cloud, pcl::KdTreeFLANN<pcl::PointXYZ>::Ptr &tree)
    {
        cloud_map_ = cloud;
        kdtree_ = tree;
        checker_->setMap(kdtree_);
    }

    void OMPLPlanner::setParams(const ros::NodeHandle &nh)
    {
        // x\y\z range以及solve_time后续可以动态配置
        nh.param("ompl_planner/secure_dis", secure_dis_, 1.0);
        nh.param("ompl_planner/x_min", x_min_, -10.0);
        nh.param("ompl_planner/x_max", x_max_, 60.0);
        nh.param("ompl_planner/y_min", y_min_, -10.0);
        nh.param("ompl_planner/y_max", y_max_, 60.0);
        nh.param("ompl_planner/z_min", z_min_, -5.0);
        nh.param("ompl_planner/z_max", z_max_, 5.0);
        nh.param("ompl_planner/solve_time", solve_time_, 5.0);
        nh.param("ompl_planner/rrt_gap", rrt_gap_, 0.3);
        nh.param("ompl_planner/dynamic_configure", dynamic_configure_, false);
        nh.param("ompl_planner/dynamic_gap", dynamic_gap_, 5.0);
        nh.param("ompl_planner/print_info", print_info_, false);
        nh.param("ompl_planner/opt_func", opt_func_, 0);
        nh.param("ompl_planner/wc", wc_, 1.0);
        nh.param("ompl_planner/wl", wl_, 10.0);

        // 设置搜索边界
        ob::RealVectorBounds bounds(3);
        bounds.setLow(0, x_min_);
        bounds.setHigh(0, x_max_);
        bounds.setLow(1, y_min_);
        bounds.setHigh(1, y_max_);
        bounds.setLow(2, z_min_);
        bounds.setHigh(2, z_max_);

        space_.reset(new ob::RealVectorStateSpace(3));
        space_->as<ob::RealVectorStateSpace>()->setBounds(bounds);

        ob::ScopedState<ob::RealVectorStateSpace> start(space_);
        ob::ScopedState<ob::RealVectorStateSpace> goal(space_);
        start->values[0] = 0;
        start->values[1] = 0;
        start->values[2] = 0;
        goal->values[0] = 0;
        goal->values[1] = 0;
        goal->values[2] = 0;

        si_.reset(new ob::SpaceInformation(space_));
        checker_.reset(new ValidityChecker(si_));
        checker_->setSecureDis(secure_dis_);
        si_->setStateValidityChecker(checker_);

        // 创建问题实例
        pdef_.reset(new ob::ProblemDefinition(si_));
        pdef_->setStartAndGoalStates(start, goal);

        // 设置优化目标
        if (opt_func_ == 0)
            pdef_->setOptimizationObjective(OMPLPlanner::getPathLengthObjWithCostToGo(si_));
        else if (opt_func_ == 1)
            pdef_->setOptimizationObjective(OMPLPlanner::getBalancedObjective(si_));
        else
            cout << "invaild opt_func parameter";
    }

    // 规划路径
    bool OMPLPlanner::plan()
    {
        // 动态配置参数
        if (dynamic_configure_)
            dynamicInitParams();

        ob::ScopedState<ob::RealVectorStateSpace> start(space_);
        start->values[0] = start_pos_[0];
        start->values[1] = start_pos_[1];
        start->values[2] = start_pos_[2];
        pdef_->clearStartStates();
        pdef_->addStartState(start);

        ob::ScopedState<ob::RealVectorStateSpace> goal(space_);
        goal->values[0] = goal_pos_[0];
        goal->values[1] = goal_pos_[1];
        goal->values[2] = goal_pos_[2];
        pdef_->clearGoal();
        pdef_->setGoalState(goal);

        // 使用RRT*算法求解
        std::shared_ptr<og::InformedRRTstar> rrt(new og::InformedRRTstar(si_));
        rrt->setRange(rrt_gap_);
        ob::PlannerPtr plan(rrt);
        plan->setProblemDefinition(pdef_);
        plan->setup();

        // 开始路径规划
        ob::PlannerStatus solved;
        try
        {
            solved = plan->solve(solve_time_);
        }
        catch (const std::exception &e)
        {
            std::cerr << e.what() << '\n';
            return false;
        }

        if (solved)
        {
            path_.reset(new og::PathGeometric(dynamic_cast<const og::PathGeometric &>(*pdef_->getSolutionPath())));
            if (print_info_)
            {
                cout << "Found solution:" << endl;
                path_->printAsMatrix(cout);
            }
            pdef_->clearSolutionPaths();
            return true;
        }
        else
            cout << "No solution found" << endl;
        return false;
    }

    vector<Vector3d> OMPLPlanner::getPath(bool smooth)
    {
        vector<Eigen::Vector3d> way_pts;
        if (smooth)
        {
            og::PathSimplifierPtr pathBSpline(new og::PathSimplifier(si_));
            pathBSpline->smoothBSpline(*path_, 3);
        }

        for (std::size_t idx = 0; idx < path_->getStateCount(); idx++)
        {
            const ob::RealVectorStateSpace::StateType *pos =
                path_->getState(idx)->as<ob::RealVectorStateSpace::StateType>();
            way_pts.push_back(Eigen::Vector3d(pos->values[0], pos->values[1], pos->values[2]));
        }
        return way_pts;
    }

    ob::OptimizationObjectivePtr OMPLPlanner::getPathLengthObjWithCostToGo(const ob::SpaceInformationPtr &si_)
    {
        ob::OptimizationObjectivePtr obj(new ob::PathLengthOptimizationObjective(si_));
        obj->setCostToGoHeuristic(&ob::goalRegionCostToGo);
        return obj;
    }

    ob::OptimizationObjectivePtr OMPLPlanner::getBalancedObjective(const ob::SpaceInformationPtr &si_)
    {
        ob::OptimizationObjectivePtr lengthObj(new ob::PathLengthOptimizationObjective(si_));
        ob::OptimizationObjectivePtr clearObj(new ClearanceObjective(si_));

        return wl_ * lengthObj + wc_ * clearObj;
    }

    void OMPLPlanner::dynamicInitParams()
    {
        // 设置搜索边界
        x_min_ = std::min(start_pos_[0], goal_pos_[0]) - dynamic_gap_;
        x_max_ = std::max(start_pos_[0], goal_pos_[0]) + dynamic_gap_;
        y_min_ = std::min(start_pos_[1], goal_pos_[1]) - dynamic_gap_;
        y_max_ = std::max(start_pos_[1], goal_pos_[1]) + dynamic_gap_;
        z_min_ = std::min(start_pos_[2], goal_pos_[2]) - dynamic_gap_;
        z_max_ = std::max(start_pos_[2], goal_pos_[2]) + dynamic_gap_;

        ob::RealVectorBounds bounds(3);
        bounds.setLow(0, x_min_);
        bounds.setHigh(0, x_max_);
        bounds.setLow(1, y_min_);
        bounds.setHigh(1, y_max_);
        bounds.setLow(2, z_min_);
        bounds.setHigh(2, z_max_);

        space_.reset(new ob::RealVectorStateSpace(3));
        space_->as<ob::RealVectorStateSpace>()->setBounds(bounds);

        si_.reset(new ob::SpaceInformation(space_));
        // checker_需要reset，因为si_是动态变化的，另外api中只有构造函数传递si的接口
        checker_.reset(new ValidityChecker(si_));
        checker_->setSecureDis(secure_dis_);
        checker_->setMap(kdtree_);
        si_->setStateValidityChecker(checker_);

        // 创建问题实例
        pdef_.reset(new ob::ProblemDefinition(si_));

        // 设置优化目标
        if (opt_func_ == 0)
            pdef_->setOptimizationObjective(OMPLPlanner::getPathLengthObjWithCostToGo(si_));
        else if (opt_func_ == 1)
            pdef_->setOptimizationObjective(OMPLPlanner::getBalancedObjective(si_));
        else
            cout << "invaild opt_func parameter";
    }
} // namespace Planner